# 调试/实验脚本
"""
代码尝试专用脚本（空壳模板）
- 用于调试、测试新功能或片段
使用 rich 替换终端输出与进度条。
"""
import time
from rich.progress import Progress, SpinnerColumn, BarColumn, TextColumn

from utils.plot_style import print_start, print_end, set_chinese_font, set_plot_style, force_chinese_font

from src.plot import prepare_border_counties_population_data
from src.gif_generator import generate_density_gif
import matplotlib as mpl
import matplotlib.pyplot as plt
from matplotlib.ticker import FuncFormatter
from pathlib import Path
import matplotlib.font_manager as fm


def run_tests():
    """在这里放置你的快速试验代码或单元级小片段。"""
    from rich.console import Console
    console = Console()
    console.rule("try.py: 运行示例测试")
    # 非破坏性测试：尝试读取 config 指定的数据文件（若不存在则打印提示）
    try:
        from src.core_utils import load_data
        try:
            df = load_data()
            console.log(f"load_data 成功，rows={len(df)} cols={len(df.columns)}")
            # 仅示例性加载数据并打印行列信息（非破坏性）
            console.log("注意：回归与示例已被移除，请参见 examples/ 目录中的示例脚本。")
        except FileNotFoundError as e:
            console.log(f"load_data 未找到文件：{e}")
        except Exception as e:
            console.log(f"load_data 出错：{e}")
    except Exception:
        console.log("src.core_utils.load_data 不可用：跳过数据加载测试")
    # 演示 rich 进度条
    with Progress(
        SpinnerColumn(style="red"),
        TextColumn("[bold]测试进度"),
        BarColumn(bar_width=None, complete_style="red"),
        TextColumn("{task.completed}/{task.total}"),
    ) as progress:
        task = progress.add_task("test", total=5)
        for _ in range(5):
            time.sleep(0.2)
            progress.advance(task)


def generate_heatmap_from_prepared_data(year: int = 2020, out_dir: str = 'figures', scope: str = 'county') -> str:
    """使用 src.plot.prepare_border_counties_population_data 准备数据并在此处完成绘图与保存。

    返回保存的相对路径字符串。
    """
    out_dir = Path(out_dir)
    out_dir.mkdir(parents=True, exist_ok=True)

    # 准备数据（不在 src.plot 中绘图/保存）
    prepared = prepare_border_counties_population_data(year=year, scope=scope)
    geo = prepared['geo']
    cmap = prepared['cmap']
    vmin = prepared['vmin']
    vmax = prepared['vmax']
    plot_figsize = prepared['figsize']
    plot_dpi = prepared['dpi']
    filename = prepared['filename']
    # filename 由 prepare 返回（已包含中文命名），无需在此追加后缀
    cb_bbox = prepared['colorbar_bbox']
    font_family = prepared['font_family']

    # 应用样式配置（如果 prepare 返回了 cfg_plot，则使用它以保持一致性）
    plot_cfg = prepared.get('cfg_plot') or {}
    set_plot_style(plot_cfg or {})

    # 仅在字体可用时加载中文字体，避免重复或不可用字体导致的警告
    if prepared.get('font_available') and prepared.get('font_family'):
        set_chinese_font(prepared.get('font_family'))
    else:
        # 建议字体不可用或未指定：静默回退到系统默认字体（避免在终端生成多余提示）
        pass

    fig, ax = plt.subplots(figsize=plot_figsize, dpi=plot_dpi)
    geo.plot(column='总人口', cmap=cmap, linewidth=0.4, edgecolor='#444444', ax=ax)
    ax.set_axis_off()

    mappable = mpl.cm.ScalarMappable(norm=mpl.colors.Normalize(vmin=vmin, vmax=vmax), cmap=cmap)
    mappable._A = []
    cax = fig.add_axes(cb_bbox)
    cb = fig.colorbar(mappable, cax=cax, orientation='horizontal')

    def per_ten_thousand(x, pos):
        # 静态图以“万人”为单位显示 -> 将原始数值除以 10000
        try:
            val = x / 10000.0
            if abs(val) >= 10:
                return f"{val:.0f}"
            return f"{val:.1f}"
        except Exception:
            return str(x)

    cb.ax.xaxis.set_major_formatter(FuncFormatter(per_ten_thousand))
    cb.ax.tick_params(labelsize=9)
    # 静态 PNG 的图例标题根据 scope 设定：
    if scope == 'city':
        cb.set_label('七普边境城市乡镇总人口数（万人）', fontsize=11)
    else:
        cb.set_label('七普边境县乡镇总人口数（万人）', fontsize=11)

    force_chinese_font(ax, font_family)
    for label in cb.ax.get_xticklabels():
        label.set_fontfamily(font_family)
    cb.ax.xaxis.label.set_fontfamily(font_family)

    fig.subplots_adjust(top=0.92, left=0.06, right=0.98, bottom=0.06)
    out_path = out_dir / filename
    bbox_opt = (prepared['cfg_plot'].get('bbox_inches') if prepared.get('cfg_plot') and prepared['cfg_plot'].get('bbox_inches') else 'tight')
    fig.savefig(out_path, dpi=plot_dpi, bbox_inches=bbox_opt)
    # 不再保存额外副本，终端只需要显示已经保存的主要 PNG 与 GIF
    repo_root = Path(__file__).resolve().parents[1]
    try:
        rel_path = out_path.relative_to(repo_root)
    except Exception:
        rel_path = out_path
    from rich.console import Console
    Console().log(f'已保存热力图到 {rel_path}')
    return str(rel_path)





if __name__ == "__main__":
    start_time = print_start('生成热力图（try.py）')
    # 直接调用 gif generator 进行演示（try.py 仅作示例调用）
    try:
        # 生成并保存按县的静态图
        generate_heatmap_from_prepared_data(year=2020, out_dir='figures', scope='county')
        # 生成并保存按市的静态图（按边境城市挑选）
        generate_heatmap_from_prepared_data(year=2020, out_dir='figures', scope='city')

        # 从配置文件读取 GIF 任务参数，允许配置覆盖默认值
        repo_root = Path(__file__).resolve().parents[1]
        cfg_path = repo_root / 'config.yaml'
        cfg = {}
        if cfg_path.exists():
            import yaml
            with open(cfg_path, 'r', encoding='utf-8') as fh:
                cfg = yaml.safe_load(fh) or {}

        # 默认的 GIF 任务列表（可在 config.yaml 中通过 add_gifs 添加/覆盖）
        default_gifs = [
            {
                'name': 'population_density',
                'years': list(range(2000, 2021)),
                'out_path': 'figures/云南_边境乡镇_人口密度均值_2000_2020.gif',
                'population_parquet': 'data/边境城市乡镇人口密度面板数据.parquet',
                'column': None,
                'cmap': None,
                'duration': 0.6,
                'verbose': True,
            },
            {
                'name': 'irri_syn',
                'years': list(range(2000, 2020)),
                'out_path': 'figures/云南_边境乡镇_IrriMap_Syn_2000_2019.gif',
                'population_parquet': 'data/边境乡镇灌溉面积占比数据.parquet',
                'column': 'IrriMap_Syn',
                'cmap': 'YlGn',
                'colorbar_label': '有效灌溉面积占比',
                'colorbar_tick_percent': True,
                'duration': 0.6,
                'verbose': True,
            }
        ]

        # 用户可以在 config.yaml 中指定 plot.gifs (覆盖默认任务列表)
        gifs_cfg = cfg.get('plot', {}).get('gifs') if cfg.get('plot') else None
        tasks = gifs_cfg or cfg.get('gifs') or default_gifs

        from rich.console import Console
        console = Console()

        for task in tasks:
            years = task.get('years')
            if isinstance(years, list):
                years_iter = range(min(years), max(years) + 1)
            else:
                years_iter = years

            out_path = task.get('out_path')
            population_parquet = task.get('population_parquet')
            column = task.get('column')
            cmap = task.get('cmap')
            duration = task.get('duration')
            verbose = task.get('verbose', False)

            try:
                gif_path, frames = generate_density_gif(
                    years=years_iter,
                    out_path=out_path,
                    population_parquet=population_parquet,
                    column=column,
                    cmap=cmap,
                    duration=duration,
                    verbose=verbose,
                )
                try:
                    rel = Path(gif_path).relative_to(repo_root)
                except Exception:
                    rel = gif_path
                console.log(f'已保存 GIF 到 {rel} (frames={frames})')
            except Exception as e:
                console.log(f'任务 {task.get("name")} 生成失败: {e}')
                import traceback
                console.log(traceback.format_exc())
    except Exception as e:
        import traceback
        from rich.console import Console
        Console().log(f'生成演示出错：{e}')
        Console().log(traceback.format_exc())
    print_end(start_time, '生成热力图（try.py）')

